package top.dprkshow.practice.graph;

import java.util.ArrayDeque;
import java.util.PriorityQueue;
import java.util.Queue;

/**
 * @author Kent HAN
 * @date 2021/9/16
 */
public class Prim {
    private boolean[] inMst;
    private PriorityQueue<Edge> pq;
    private Queue<Edge> mst;

    public Prim(WeightedUgraph g) {
        inMst = new boolean[g.getTotalVertices()];
        pq = new PriorityQueue<>();
        mst = new ArrayDeque<>();
    }

    public void init(WeightedUgraph g) {
        visit(g, 0);
        while(!pq.isEmpty() && mst.size() < g.getTotalVertices()) {
            Edge e = pq.poll();
            int v = e.either(), w = e.other(v);
            if (inMst[v] && inMst[w]) {
                continue;
            }
            mst.add(e);
            if (!inMst[v]) {
                visit(g, v);
            }
            if (!inMst[w]) {
                visit(g, w);
            }
        }
    }

    public Iterable<Edge> mst() {
        return mst;
    }

    private void visit(WeightedUgraph g, int v) {
        inMst[v] = true;
        for (Edge e : g.adj(v)) {
            if (!inMst[e.other(v)]) {
                pq.add(e);
            }
        }
    }
}
